{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7ff2c269",
   "metadata": {},
   "source": [
    "# Sentence Embedding Optimizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "839c4a87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# My OpenAI Key\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"INSERT OPENAI KEY\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "40cf0773",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "fa34cd83",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index import download_loader\n",
    "\n",
    "WikipediaReader = download_loader(\"WikipediaReader\")\n",
    "\n",
    "loader = WikipediaReader()\n",
    "documents = loader.load_data(pages=[\"Berlin\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f59e6c18",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'llama_index.readers.schema.base.Document'>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [build_index_from_documents] Total LLM token usage: 0 tokens\n",
      "INFO:root:> [build_index_from_documents] Total embedding token usage: 18390 tokens\n"
     ]
    }
   ],
   "source": [
    "from llama_index import VectorStoreIndex\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "827ada33",
   "metadata": {},
   "source": [
    "Compare query with and without optimization for LLM token usage, Embedding Model usage on query, Embedding model usage for optimizer, and total time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a04e4535",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Without optimization\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [query] Total LLM token usage: 3545 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 7 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total time elapsed: 2.8928110599517822\n",
      "Answer: \n",
      "The population of Berlin in 1949 was approximately 2.2 million inhabitants. After the fall of the Berlin Wall in 1989, the population of Berlin increased to approximately 3.7 million inhabitants.\n",
      "\n",
      "With optimization\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [optimize] Total embedding token usage: 7 tokens\n",
      "INFO:root:> [query] Total LLM token usage: 1779 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 7 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total time elapsed: 2.346346139907837\n",
      "Answer: \n",
      "The population of Berlin is around 4.5 million.\n",
      "Alternate optimization cutoff\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:> [optimize] Total embedding token usage: 7 tokens\n",
      "INFO:root:> [query] Total LLM token usage: 3215 tokens\n",
      "INFO:root:> [query] Total embedding token usage: 7 tokens\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total time elapsed: 2.101111888885498\n",
      "Answer: \n",
      "The population of Berlin is around 4.5 million.\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "from llama_index import VectorStoreIndex\n",
    "from llama_index.indices.postprocessor import SentenceEmbeddingOptimizer\n",
    "\n",
    "print(\"Without optimization\")\n",
    "start_time = time.time()\n",
    "query_engine = index.as_query_engine()\n",
    "res = query_engine.query(\"What is the population of Berlin?\")\n",
    "end_time = time.time()\n",
    "print(\"Total time elapsed: {}\".format(end_time - start_time))\n",
    "print(\"Answer: {}\".format(res))\n",
    "\n",
    "print(\"With optimization\")\n",
    "start_time = time.time()\n",
    "query_engine = index.as_query_engine(\n",
    "    node_postprocessors=[SentenceEmbeddingOptimizer(percentile_cutoff=0.5)]\n",
    ")\n",
    "res = query_engine.query(\"What is the population of Berlin?\")\n",
    "end_time = time.time()\n",
    "print(\"Total time elapsed: {}\".format(end_time - start_time))\n",
    "print(\"Answer: {}\".format(res))\n",
    "\n",
    "print(\"Alternate optimization cutoff\")\n",
    "start_time = time.time()\n",
    "query_engine = index.as_query_engine(\n",
    "    node_postprocessors=[SentenceEmbeddingOptimizer(threshold_cutoff=0.7)]\n",
    ")\n",
    "res = query_engine.query(\"What is the population of Berlin?\")\n",
    "end_time = time.time()\n",
    "print(\"Total time elapsed: {}\".format(end_time - start_time))\n",
    "print(\"Answer: {}\".format(res))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
